#!/usr/bin/env python 
#  thirdorder, help compute anharmonic IFCs from minimal sets of displacements
#  Copyright (C) 2012-2018 Wu Li <wu.li.phys2011@gmail.com>
#  Copyright (C) 2012-2018 Jesús Carrete Montaña <jcarrete@gmail.com>
#  Copyright (C) 2012-2018 Natalio Mingo Bisquert <natalio.mingo@cea.fr>
#  Copyright (C) 2014-2018 Antti J. Karttunen <antti.j.karttunen@iki.fi>
#  Copyright (C) 2016-2018 Genadi Naydenov <gan503@york.ac.uk>
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
import os
from os import listdir,system
import sys
import os.path
from ase.calculators.lammpsrun import LAMMPS

try:
    import cStringIO as StringIO
except ImportError:
    try:
        import StringIO
    except ImportError:
        import io as StringIO
try:
    import hashlib
    hashes = True
except ImportError:
    hashes = False

import thirdorder_core
from thirdorder_common import *
from ase.io import read
from ase import Atoms
from ase.io.trajectory import TrajectoryWriter
from ase.calculators.singlepoint import SinglePointCalculator
from irff.md.lammps import writeLammpsData,writeLammpsIn,lammpstraj_to_ase
from irff.md.lammps import get_lammps_forces 

def read_POSCAR(directory):
    """
    Return all the relevant information contained in a POSCAR file.
    """
    with dir_context(directory):
        nruter = dict()
        nruter["lattvec"] = np.empty((3, 3))
        f = open("POSCAR.0", "r")
        firstline = next(f)
        factor = .1 * float(next(f).strip())
        for i in range(3):
            nruter["lattvec"][:, i] = [float(j) for j in next(f).split()]
        nruter["lattvec"] *= factor
        line = next(f)
        fields = next(f).split()
        old = False
        try:
            int(fields[0])
        except ValueError:
            old = True
        if old:
            nruter["elements"] = firstline.split()
            nruter["numbers"] = np.array([int(i) for i in line.split()])
            typeline = "".join(fields)
        else:
            nruter["elements"] = line.split()
            nruter["numbers"] = np.array(
                [int(i) for i in fields], dtype=np.intc)
            typeline = next(f)
        natoms = nruter["numbers"].sum()
        nruter["positions"] = np.empty((3, natoms))
        for i in range(natoms):
            nruter["positions"][:, i] = [float(j) for j in next(f).split()]
        f.close()
    nruter["types"] = []
    for i in range(len(nruter["numbers"])):
        nruter["types"] += [i] * nruter["numbers"][i]
    if typeline[0] == "C":
        nruter["positions"] = sp.linalg.solve(nruter["lattvec"],
                                              nruter["positions"] * factor)
    return nruter


def write_POSCAR(poscar, filename):
    """
    Write the contents of poscar to filename.
    """
    global hashes
    f = StringIO.StringIO()
    f.write("1.0\n")
    for i in range(3):
        f.write("{0[0]:>20.15f} {0[1]:>20.15f} {0[2]:>20.15f}\n".format((
            poscar["lattvec"][:, i] * 10.).tolist()))
    f.write("{0}\n".format(" ".join(poscar["elements"])))
    f.write("{0}\n".format(" ".join([str(i) for i in poscar["numbers"]])))
    f.write("Direct\n")
    for i in range(poscar["positions"].shape[1]):
        f.write("{0[0]:>20.15f} {0[1]:>20.15f} {0[2]:>20.15f}\n".format(
            poscar["positions"][:, i].tolist()))
    if hashes:
        toencode = f.getvalue()
        if sys.hexversion >= 0x3000000:
            toencode = toencode.encode("utf-8")
        header = hashlib.sha1(toencode).hexdigest()
    else:
        header = filename
    with open(filename, "w") as finalf:
        finalf.write("{0}\n".format(header))
        finalf.write(f.getvalue())
    f.close()


def normalize_SPOSCAR(sposcar):
    """
    Rearrange sposcar, as generated by gen_SPOSCAR, so that it is in
    valid VASP order, and return the result.
    """
    nruter = copy.deepcopy(sposcar)
    # Order used internally (from most to least significant):
    # k,j,i,iat For VASP, iat must be the most significant index,
    # i.e., atoms of the same element must go together.
    indices = np.array(range(nruter["positions"].shape[1])).reshape(
        (sposcar["nc"], sposcar["nb"], sposcar["na"], -1))
    indices = np.rollaxis(indices, 3, 0).flatten().tolist()
    nruter["positions"] = nruter["positions"][:, indices]
    nruter["types"].sort()
    return nruter


def read_forces(filename):
    """
    Read a set of forces on atoms from filename, presumably in
    vasprun.xml format.
    """
    calculation = ElementTree.parse(filename).getroot().find("calculation")
    for a in calculation.findall("varray"):
        if a.attrib["name"] == "forces":
            break
    nruter = []
    for i in a.getchildren():
        nruter.append([float(j) for j in i.text.split()])
    nruter = np.array(nruter, dtype=np.double)
    return nruter


def build_unpermutation(sposcar):
    """
    Return a list of integers mapping the atoms in the normalized
    version of sposcar to their original indices.
    """
    indices = np.array(range(sposcar["positions"].shape[1])).reshape(
        (sposcar["nc"], sposcar["nb"], sposcar["na"], -1))
    indices = np.rollaxis(indices, 3, 0).flatten()
    return indices.argsort().tolist()


if __name__ == "__main__":
    ''' thirdorder_lammps.py sow|reap na nb nc cutoff [nm/-integer] 
        thirdorder_lammps.py sow 4 4 1 -3  # create input file
        thirdorder_lammps.py sow 4 4 1 10  # create input file
        thirdorder_lammps.py sow 4 4 1 -3  # create force constant 3rd 
    '''
    if len(sys.argv) < 6 or sys.argv[1] not in ("sow", "reap"):
        sys.exit("Usage: {0} sow|reap na nb nc cutoff[nm/-integer] [calculator]".format(
            sys.argv[0]))
    action = sys.argv[1]
    if len(sys.argv) >6 :
       pair_style= sys.argv[6]
    if len(sys.argv) >7 :
       pair_coeff= sys.argv[7]
     
    na, nb, nc = [int(i) for i in sys.argv[2:5]]

    if min(na, nb, nc) < 1:
        sys.exit("Error: na, nb and nc must be positive integers")
    if sys.argv[5][0] == "-":
        try:
            nneigh = -int(sys.argv[5])
        except ValueError:
            sys.exit("Error: invalid cutoff")
        if nneigh == 0:
            sys.exit("Error: invalid cutoff")
    else:
        nneigh = None
        try:
            frange = float(sys.argv[5])
        except ValueError:
            sys.exit("Error: invalid cutoff")
        if frange == 0.:
            sys.exit("Error: invalid cutoff")      
#   - end - parse argument options - end - parse argument options - end - parse argument options

    print("Reading POSCAR.0")
    poscar = read_POSCAR(".")
    natoms = len(poscar["types"])
    print("Analyzing the symmetries")
    symops = thirdorder_core.SymmetryOperations(
        poscar["lattvec"], poscar["types"], poscar["positions"].T, SYMPREC)
    print("- Symmetry group {0} detected".format(symops.symbol))
    print("- {0} symmetry operations".format(symops.translations.shape[0]))
    print("Creating the supercell")
    sposcar = gen_SPOSCAR(poscar, na, nb, nc)
    ntot = natoms * na * nb * nc
    print("Computing all distances in the supercell")
    dmin, nequi, shifts = calc_dists(sposcar)
    if nneigh != None:
       frange = calc_frange(poscar, sposcar, nneigh, dmin)
       print("- Automatic cutoff: {0} nm".format(frange))
    else:
       print("- User-defined cutoff: {0} nm".format(frange))

    print("Looking for an irreducible set of third-order IFCs")
    wedge = thirdorder_core.Wedge(poscar, sposcar, symops, dmin, nequi, shifts,
                                  frange)
    print("- {0} triplet equivalence classes found".format(wedge.nlist))
    list4 = wedge.build_list4()
    nirred = len(list4)
    nruns = 4 * nirred
    print("- {0} DFT runs are needed".format(nruns))
    if action == "sow":
        print(sowblock)
        print("Writing undisplaced coordinates to 3RD.SPOSCAR")
        write_POSCAR(normalize_SPOSCAR(sposcar), "3RD.SPOSCAR")
        width = len(str(4 * (len(list4) + 1)))
        namepattern = "3RD.POSCAR.{{0:0{0}d}}".format(width)
        print("Writing displaced coordinates to 3RD.POSCAR.*")
        for i, e in enumerate(list4):
            for n in range(4):
                isign = (-1)**(n // 2)
                jsign = -(-1)**(n % 2)
                # Start numbering the files at 1 for aesthetic
                # reasons.
                number = nirred * n + i + 1
                dsposcar = normalize_SPOSCAR(
                    move_two_atoms(sposcar, e[1], e[3], isign * H, e[0], e[2],
                                   jsign * H))
                filename = namepattern.format(number)
                write_POSCAR(dsposcar, filename)
    else: 
#     --- generating forces - generating forces - generating forces - generating forces ---
        # try: 
        #    system('mv POSCAR POSCAR.0')
        # except:
        #    print(' ')
        print(reapblock)                        # compute force constant
        # print("XML ElementTree implementation: {0}".format(xmllib))
        # print("Waiting for a list of vasprun.xml files on stdin")
        cars = listdir()
        filelist = []
        ind_     = []
        for car in cars:
           if car.find('3RD.POSCAR')>=0:
              # system('mv '+ car + 'POSCAR')
              ind_.append(int(car.split('.')[2]))
              filelist.append(car)

        ind = np.argsort(ind_)
        filelist = np.array(filelist)[ind]
        # atoms = read('POSCAR')
        # atoms = get_gulp_forces([atoms])
        nfiles = len(filelist)
        print("- {0} filenames read".format(nfiles))

        if nfiles != nruns:
           sys.exit("Error: {0} filenames were expected".format(nruns))
        for i in filelist:
            if not os.path.isfile(i):
               sys.exit("Error: {0} is not a regular file".format(i))
        print("Reading the forces")

        p = build_unpermutation(sposcar)
        
        forces = []
        for i in filelist:
            # forces.append(read_forces(i)[p, :])
            system('cp '+ i + ' POSCAR')
            atoms = read('POSCAR')
            #atoms = get_gulp_forces([atoms])
            # if len(sys.argv) > 7:
            #    atoms = get_lammps_forces(atoms,ffield='ffield',pair_style=pair_style,
            #                              pair_coeff=pair_coeff) #
            # elif len(sys.argv) > 6:
            #    atoms = get_lammps_forces(atoms,ffield='ffield',pair_style=pair_style) #
            # else:
            #    atoms = get_lammps_forces(atoms,ffield='ffield') 
            files = ['Carbon_GAP_20_potential/Carbon_GAP_20.xml',
             'Carbon_GAP_20_potential/Carbon_GAP_20.xml.sparseX.GAP_2020_4_27_60_2_50_5_4361',
             'Carbon_GAP_20_potential/Carbon_GAP_20.xml.sparseX.GAP_2020_4_27_60_2_50_5_4362',
             'Carbon_GAP_20_potential/Carbon_GAP_20.xml.sparseX.GAP_2020_4_27_60_2_50_5_4363']
            os.environ['ASE_LAMMPSRUN_COMMAND'] = 'mpirun -n 4 lammps'
            lammps = LAMMPS(files=files)
            lammps.set(pair_style='quip')
            lammps.set(pair_coeff=['* * Carbon_GAP_20_potential/Carbon_GAP_20.xml \"\" 6'])
            lammps.set(tmp_dir='./')
            lammps.set(keep_alive=False)
            # lammps.set(keep_tmp_files=True)
            atoms.calc = lammps
            forces.append(atoms.get_forces()[p, :])
            print("- {0} read successfully".format(i))
            res = forces[-1].mean(axis=0)
            print("- \t Average force:")
            print("- \t {0} eV/(A * atom)".format(res))
        print("Computing an irreducible set of anharmonic force constants")
        
        phipart = np.zeros((3, nirred, ntot))
        for i, e in enumerate(list4):
            for n in range(4):
                isign = (-1)**(n // 2)
                jsign = -(-1)**(n % 2)
                number = nirred * n + i
                phipart[:, i, :] -= isign * jsign * forces[number].T
        phipart /= (400. * H * H)
        print("Reconstructing the full array")
        phifull = thirdorder_core.reconstruct_ifcs(phipart, wedge, list4,
                                                   poscar, sposcar)
        print("Writing the constants to FORCE_CONSTANTS_3RD")
        write_ifcs(phifull, poscar, sposcar, dmin, nequi, shifts, frange,
                   "FORCE_CONSTANTS_3RD")
    print(doneblock)
